// Copyright (c) 2022 Presto Labs Pte. Ltd.
// Author: jhkim

#pragma once
#include <algorithm>
#include <map>
#include <vector>
#include <memory>
#include <string>
#include <utility>
#include <unordered_set>

#include "presto/quant/feature/focus_rule/focus_indexer.h"

namespace coin2::strategy::linear_model {

using ::presto::quant::feature::GetNormBase;
using ::presto::quant::feature::GetQuote;
using ::presto::quant::feature::IsTrueBasisSymbol;

struct LinearModelWrapper {
  int64_t lm_feed_id = 0L;
  std::unique_ptr<LinearModel> linear_model;
  LinearModelDelegateProto delegate;
};

inline void erase_duplicate(std::vector<::presto::Mep>* meps) {
  std::sort(
      meps->begin(),
      meps->end(),
      [] (const auto& mep1, const auto& mep2) -> bool {
        return std::make_tuple(mep1.market_type, mep1.exchange_name, mep1.product_name)
            < std::make_tuple(mep2.market_type, mep2.exchange_name, mep2.product_name);
      });
  meps->erase(std::unique(
      meps->begin(),
      meps->end(),
      [] (const auto& mep1, const auto& mep2) -> bool {
        return std::make_tuple(mep1.market_type, mep1.exchange_name, mep1.product_name)
            == std::make_tuple(mep2.market_type, mep2.exchange_name, mep2.product_name);
      }),
      meps->end());
}

inline void add_symbols_reserve(
    std::vector<::presto::Mep>* meps,
    const LmFocusConfig* fcss) {
  MarketType market_type = coin::proto::UNKNOWN_MARKET_TYPE;
  ExchangeType exchange_type = coin::proto::UNKNOWN_EXCHANGE_TYPE;
  for (const auto& sr: fcss->symbols_reserve()) {
    if (!MarketType_Parse(sr.market_type(), &market_type)) {
      LOG(ERROR) << "Invalid market_type on symbols_reserve: " << sr.market_type();
      continue;
    } 
    if (!ExchangeType_Parse(sr.exchange_name(), &exchange_type)) {
      LOG(ERROR) << "Invalid exchange_name on symbols_reserve: " << sr.exchange_name();
      continue;
    } 
    meps->emplace_back(presto::Mep{sr.market_type(), sr.exchange_name(), sr.product_name()});
  }
}

inline void UnSubscribeUntradeRefFromFeedMgr(
    coin2::feed::mgr::FeedSubsystemManager* feed_mgr,
    const std::set<::presto::Mep>& trade_fcs_list,
    const std::set<::presto::Mep>& untrade_but_subscribe_fcs_list) {
  std::set<std::string> ban_bases;
  for (auto product : feed_mgr->GetProducts()) {
    auto mep_fs = presto::Mep {
          MarketType_Name(product->market()),
          ExchangeType_Name(product->exchange()),
          product->relative_norm()
      };
    if (untrade_but_subscribe_fcs_list.count(mep_fs) > 0) {
      std::string base = GetNormBase(product->base().currency());
      ban_bases.insert(base);
    }
  }

  for (auto&& feed_subsystem : feed_mgr->GetAllFeedSubsystems()) {
    std::vector<const IProduct*> ban_products;
    for (const auto product : feed_subsystem->GetProducts()) {
      auto mep_fs = presto::Mep {
          MarketType_Name(product->market()),
          ExchangeType_Name(product->exchange()),
          product->relative_norm()
      };

      if (trade_fcs_list.count(mep_fs) == 0
          && untrade_but_subscribe_fcs_list.count(mep_fs) == 0
          && ban_bases.count(GetNormBase(product->base().currency())) > 0) {
        // it means ref
        ban_products.emplace_back(product);
      }
    }
    if (!ban_products.empty()) {
      CHECK(feed_subsystem->BanProduct(ban_products));
      for (auto ban_product : ban_products) {
        LOG(INFO) << "Ban Ref: " << ban_product->relative_norm();
      }
    }
  }
}

inline std::vector<::presto::Mep> UnSubscribeFromFeedMgr(
    coin2::feed::mgr::FeedSubsystemManager* feed_mgr,
    const std::vector<::presto::Mep>& allow_meps,
    std::vector<::presto::Mep>* clear_meps_fcsref = nullptr,
    bool unsubscribe = false,
    bool require_quote_matching = false) {
  std::vector<::presto::Mep> allow_meps_fcsref;
  std::unordered_set<std::string> allow_product_names;
  for (const auto& allow_mep : allow_meps) {
    allow_product_names.insert(allow_mep.product_name);
  }
  // just ban all the bases from feed
  std::unordered_set<std::string> allow_bases;
  for (auto product : feed_mgr->GetProducts()) {
    if (allow_product_names.count(product->relative_norm()) > 0) {
      std::string base = GetNormBase(product->base().currency());
      std::string quote = (require_quote_matching)? GetQuote(product->relative_norm()) : "";
      allow_bases.insert(base + "-" + quote);
    }
  }

  // default is 0
  auto env_feed_noban = getenv("FEED_NOBAN_TOP_VU");
  bool noban = (env_feed_noban && std::string(env_feed_noban) == std::string("1"));
  if (!unsubscribe) {
    noban = true;
  }

  for (auto&& feed_subsystem : feed_mgr->GetAllFeedSubsystems()) {
    std::vector<const IProduct*> ban_products;
    for (auto product : feed_subsystem->GetProducts()) {
      std::string base = GetNormBase(product->base().currency());
      std::string quote = (require_quote_matching)? GetQuote(product->relative_norm()) : "";
      if (allow_bases.count(base + "-" + quote) == 0) {
        if (!noban) {
          ban_products.push_back(product);
        }
        if (clear_meps_fcsref) {
          clear_meps_fcsref->push_back(
              presto::Mep{
                  MarketType_Name(product->market()),
                  ExchangeType_Name(product->exchange()),
                  product->relative_norm()
                  });
        }
      } else {
        LOG(INFO) << "NotUnsubscribed " << product->relative_norm();
        allow_meps_fcsref.push_back(presto::Mep{
            MarketType_Name(product->market()),
            ExchangeType_Name(product->exchange()),
            product->relative_norm()
            });
      }
    }
    if (!ban_products.empty()) {
      CHECK(feed_subsystem->BanProduct(ban_products));
      for (auto ban_product : ban_products) {
        LOG(INFO) << "Unsubscribed " << ban_product->relative_norm();
      }
    }
  }
  if (unsubscribe) {
    LOG(INFO) << "# Subscribing: " << allow_meps_fcsref.size();
  }
  erase_duplicate(&allow_meps_fcsref);
  if (clear_meps_fcsref) erase_duplicate(clear_meps_fcsref);
  return allow_meps_fcsref;
}

inline presto::quant::feature::linear_model::LinearAlphaProto GetLaProtoFromZip(
    const std::string& profile_zip,
    const std::string& linear_alpha_filename) {
  return presto::util::read_proto<presto::quant::feature::linear_model::LinearAlphaProto>(
      profile_zip,
      linear_alpha_filename);
}

inline FeatureSpecProto GetFSpecProtoFromZip(
    const std::string& profile_zip,
    const std::string& linear_alpha_filename) {
  auto la_proto = GetLaProtoFromZip(profile_zip, linear_alpha_filename);
  return presto::util::read_proto<FeatureSpecProto>(
      profile_zip, la_proto.feature_filename());
}

inline FeatureSpecProto GetFSpecProto(
    const std::string& feature_info_filename) {
  FeatureSpecProto fspec_pb;
  std::ifstream in(feature_info_filename);
  CHECK_THROW(in.is_open());
  std::string str((std::istreambuf_iterator<char>(in)), std::istreambuf_iterator<char>());
  boost::replace_all(str, "y-gen", "y_gen");
  CHECK(google::protobuf::util::JsonStringToMessage(str, &fspec_pb).ok());
  return fspec_pb;
}

inline std::optional<presto::Mep> GetFocusMep(
    const FeatureSpecProto& fspec_pb) {
  if (fspec_pb.focus_symbols_size() == 0) {
    return std::nullopt;
  }
  std::optional<presto::Mep> ret(presto::Mep{});

  auto exch = fspec_pb.focus_symbols(0).exchange();
  if (exch.find("Futures") != std::string::npos) {
    ret->market_type = "Futures";
  } else {
    ret->market_type = "Spot";
  }
  boost::replace_all(exch, "Futures", "");
  ret->exchange_name = exch;
  ret->product_name = fspec_pb.focus_symbols(0).product_name();
  return ret;
}

inline bool GetFocusSymbolsVuRankGeLt(
    const FeatureSpecProto& fspec_pb,
    int64_t* rank_ge,
    int64_t* rank_lt) {
  std::string expr = fspec_pb.focus_symbols_vu_rank();
  if (expr.empty()) return false;
  if (expr == "all") {
    *rank_ge = 0L;
    *rank_lt = 10000L;
  } else if (expr.find("-") == std::string::npos) {
    *rank_ge = 0L;
    *rank_lt = std::stoll(expr);
  } else {
    std::vector<std::string> splits;
    SplitString(expr, '-', &splits);
    CHECK_EQ(splits.size(), 2) << expr;
    *rank_ge = std::stoll(splits[0]);
    *rank_lt = std::stoll(splits[1]);
  }
  return true;
}

inline void AddBanMeps(
    const SymbolManager* symbol_manager,
    const std::vector<presto::Mep>& meps) {
  for (const auto& mep : meps) {
    symbol_manager->AddBanMep(mep);
  }
}

inline int64_t GetMatchingTs(const VolumeInfo& volume_info, int64_t timestamp_end) {
  CHECK(volume_info.volume_info_by_exchs_size() > 0) << "no volume_info_by_exchs";
  int64_t matching_ts = 0L;
  for (const auto& volume_info_by_exch : volume_info.volume_info_by_exchs()) {
    if (volume_info_by_exch.exch_idx() == 0) { // only for focus
      for (const auto& volume_info_i : volume_info_by_exch.volume_infos()) {
        if (matching_ts <= volume_info_i.timestamp() && volume_info_i.timestamp() <= timestamp_end) {
          matching_ts = volume_info_i.timestamp();
        }
      }
    }
  }
  return matching_ts;
}

inline std::unique_ptr<IProduct> CreateProductFromMEPStrings(
    const std::string& market_type_str,
    const std::string& exchange_name_str,
    const std::string& product_name,
    int64_t timestamp_override) {
  MarketType market_type = coin::proto::UNKNOWN_MARKET_TYPE;
  ExchangeType exchange_type = coin::proto::UNKNOWN_EXCHANGE_TYPE;
  CHECK(MarketType_Parse(market_type_str, &market_type)) << market_type_str;
  CHECK(ExchangeType_Parse(exchange_name_str, &exchange_type)) << exchange_name_str;
  return CreateProductFromNormString(market_type, exchange_type, product_name, timestamp_override);
}

inline std::vector<presto::Mep> GetSymbolsRank(
    const VolumeInfo& volume_info,
    int64_t symbol_universe_rank_ge,
    int64_t symbol_universe_rank_lt,
    int64_t timestamp_end,
    bool add_basis_coins_forcibly,
    std::vector<::presto::Mep>* symbols_non_sel = nullptr,
    std::vector<::presto::Mep>* symbols_exclude = nullptr,
    bool require_quote_matching = false,
    bool exclude_noref_symbols = true
  ) {
  CHECK(volume_info.volume_info_by_exchs_size() > 0) << "no volume_info_by_exchs";
  LOG(INFO) << "require_quote_matching=" << require_quote_matching;
  LOG(INFO) << "exclude_noref_symbols=" << exclude_noref_symbols;
  int64_t matching_ts = GetMatchingTs(volume_info, timestamp_end);
  std::map<std::pair<double, int64_t>, presto::Mep> volume_symbols_focus;
  std::unordered_map<int, std::unordered_map<std::string, bool>> ref_bases;
  for (const auto& volume_info_by_exch : volume_info.volume_info_by_exchs()) {
    int exch_idx = volume_info_by_exch.exch_idx();
    for (int i = 0; i < volume_info_by_exch.volume_infos_size(); ++i) {
      const auto& volume_info_i = volume_info_by_exch.volume_infos(i);
      if (matching_ts == volume_info_i.timestamp()) {
        if (exch_idx == 0) { // focus exchange
          auto mep = presto::Mep{
              volume_info_by_exch.market_type(),
              volume_info_by_exch.exchange_name(),
              volume_info_i.symbol()
          };
          volume_symbols_focus[std::make_pair(volume_info_i.volume(), i)] = mep;
        } else { // ref exchange
          if (ref_bases.count(exch_idx) == 0) {
            ref_bases[exch_idx] = std::unordered_map<std::string, bool>();
          }
          auto product = CreateProductFromMEPStrings(
              volume_info_by_exch.market_type(),
              volume_info_by_exch.exchange_name(),
              volume_info_i.symbol(),
              timestamp_end);
          std::string ref_base = GetNormBase(product->base().currency());
          std::string ref_quote = (require_quote_matching)? GetQuote(product->relative_norm()) : "";
          ref_bases[exch_idx][ref_base + "-" + ref_quote] = true;
        }
      }
    }
  }
  int curr_rank = 0;
  std::vector<presto::Mep> symbols_sel;
  for (auto it = volume_symbols_focus.rbegin(); it != volume_symbols_focus.rend(); ++it) {
    const auto& mep = it->second;
    auto product = CreateProductFromMEPStrings(
        mep.market_type, mep.exchange_name, mep.product_name, timestamp_end);
    std::string symbol_base = GetNormBase(product->base().currency());
    std::string symbol_quote = (require_quote_matching)? GetQuote(product->relative_norm()) : "";
    int ref_found_cnt = 0;
    for (const auto& [exch_idx, base_map] : ref_bases) {
      if (base_map.count(symbol_base + "-" + symbol_quote) > 0) {
        ref_found_cnt++;
      }
    }
    const double volume = (it->first).first;
    bool is_valid_symbol = false;
    if (ref_found_cnt < ref_bases.size() && exclude_noref_symbols) { // skip
      LOG(INFO) << "Skip Due to lack of ref: " << mep.product_name;
      if (symbols_exclude) {
        // Sometimes we need to clear. so add to exclude list to subscribe.
        symbols_exclude->push_back(mep);
      }
    } else {
      if (curr_rank >= symbol_universe_rank_ge && curr_rank < symbol_universe_rank_lt && volume > 0) {
        is_valid_symbol = true;
      } else if (add_basis_coins_forcibly && IsTrueBasisSymbol(mep.product_name)) {
        is_valid_symbol = true;
      }
      ++curr_rank;
    }

    // extract symbols_exclude
    if (volume <= 0.0 && symbols_exclude) {
      symbols_exclude->push_back(mep);
    }

    if (is_valid_symbol) {
      symbols_sel.push_back(mep);
      LOG(INFO) << mep.GetSymbolString() << " added, curr_rank=" << curr_rank;
    } else if (symbols_non_sel) { // for invalid symbol
      symbols_non_sel->push_back(mep);
    }
  }
  return symbols_sel;
}

inline std::vector<presto::Mep> BanFocusRankAndGetAllowlist(
    coin2::feed::mgr::FeedSubsystemManager* feed_mgr,
    const VolumeInfo& volume_info,
    int64_t timestamp_override,
    int64_t rank_ge,
    int64_t rank_lt,
    std::vector<presto::Mep>* clearlist_fcsref = nullptr,
    LmFocusConfig* fcss = nullptr,
    bool require_quote_matching = false,
    bool exclude_noref_symbols = true) {
  feed_mgr->DebugPrint();
  bool add_basis_coins_forcibly = true;
  int64_t live_sub_rank_lt = rank_lt;
  if (fcss && fcss->has_enable_clear_all()) {
    int clear_version = fcss->enable_clear_all();
    LOG(INFO) << "Enable Clear ver: " << clear_version;
    if (clear_version >= 1) {
      // only btc, eth
      rank_ge = rank_lt = -1;
    }
  }

  std::vector<presto::Mep> banlist_fcs, exclude_fcs;
  std::vector<presto::Mep> allowlist_fcs = GetSymbolsRank(
      volume_info,
      rank_ge,
      rank_lt,
      timestamp_override,
      add_basis_coins_forcibly,
      &banlist_fcs,
      &exclude_fcs,
      require_quote_matching,
      exclude_noref_symbols);

  std::vector<presto::Mep> subscribelist_fcs = allowlist_fcs;

  if (feed_mgr->HasLive()) {
    subscribelist_fcs = GetSymbolsRank(
        volume_info,
        0,
        live_sub_rank_lt + 80,
        timestamp_override,
        add_basis_coins_forcibly,
        nullptr,
        nullptr,
        require_quote_matching,
        exclude_noref_symbols);
    LOG(INFO) << "Subscribing " << 0 << " ~ " << live_sub_rank_lt + 80;
    LOG(INFO) << "[Symbols Exclude] - Additional Subscribe in Live to clear pos";
    for (auto ls : exclude_fcs) {
      LOG(INFO) << "## " << ls.GetSymbolString();
      subscribelist_fcs.push_back(ls);
    }
  } else {
    LOG(INFO) << "Subscribing " << rank_ge << " ~ " << rank_lt;
  }

  erase_duplicate(&subscribelist_fcs);
  if (fcss) {
    add_symbols_reserve(&subscribelist_fcs, fcss);
  }

  // Unsubscribe unnecessary products.
  LOG(INFO) << allowlist_fcs.size() << " " << subscribelist_fcs.size();
  std::set<presto::Mep> trade_fcs_list(allowlist_fcs.begin(), allowlist_fcs.end());
  std::set<presto::Mep> subscribe_fcs_list(subscribelist_fcs.begin(), subscribelist_fcs.end());
  std::set<presto::Mep> untrade_but_subscribe_fcs_list;
  std::set_difference(
      subscribe_fcs_list.begin(), subscribe_fcs_list.end(),
      trade_fcs_list.begin(), trade_fcs_list.end(),
      inserter(untrade_but_subscribe_fcs_list, untrade_but_subscribe_fcs_list.begin())
  );
  UnSubscribeFromFeedMgr(
      feed_mgr,
      subscribelist_fcs,
      clearlist_fcsref,
      true,
      require_quote_matching);
  UnSubscribeUntradeRefFromFeedMgr(
      feed_mgr,
      trade_fcs_list,
      untrade_but_subscribe_fcs_list);

  std::vector<presto::Mep> allowlist_fcsref = UnSubscribeFromFeedMgr(
      feed_mgr,
      allowlist_fcs,
      clearlist_fcsref,
      false,
      require_quote_matching);

  if (fcss) {
    // when the trade univ is a subset of the feature univ
    if (fcss->has_symbols_universe_rank()) {
      CHECK_GE(fcss->symbols_universe_rank(), rank_ge);
      CHECK_LE(fcss->symbols_universe_rank(), rank_lt);
      banlist_fcs.clear();
      allowlist_fcs = GetSymbolsRank(
          volume_info,
          rank_ge,
          fcss->symbols_universe_rank(),
          timestamp_override,
          false,
          &banlist_fcs,
          nullptr,
          require_quote_matching,
          exclude_noref_symbols);
    }
    fcss->clear_symbols();
    for (const auto& mep : banlist_fcs) {
      LOG(INFO) << "CLEAR: " << mep.product_name;
      fcss->add_symbols_clear(mep.product_name);
    }
    int clear_version = fcss->enable_clear_all();
    for (const auto& mep : allowlist_fcs) {
      if (clear_version == 1) {
        // Force Clear
        LOG(INFO) << "CLEAR: " << mep.product_name;
        fcss->add_symbols_clear(mep.product_name);
      } else {
        LOG(INFO) << "Focus: " << mep.product_name;
        fcss->add_symbols(mep.product_name);
      }
    }
  }
  return allowlist_fcsref;
}

}  // namespace coin2::strategy::linear_model

